package com.study.kafka.custom.interceptor;

import org.apache.kafka.clients.consumer.ConsumerInterceptor;
import org.apache.kafka.clients.consumer.ConsumerRecord;
import org.apache.kafka.clients.consumer.ConsumerRecords;
import org.apache.kafka.common.TopicPartition;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * 消费者拦截器
 * 使用：props.put(ConsumerConfig.INTERCEPTOR_CLASSES_CONFIG, DemoConsumerInterceptor.class.getName());
 *
 * @author: stephen.shen
 * @create: 2019-06-18 16:56
 */
public class DemoConsumerInterceptor implements ConsumerInterceptor<String, String> {

    private static final long EXPIRE_INTERVAL = 10 * 1000;

    /**
     * KafkaConsumer会在在poll()之前调用
     *
     * @param records
     * @return
     */
    @Override
    public ConsumerRecords<String, String> onConsume(ConsumerRecords<String, String> records) {
        // 实现丢弃消息，时间戳与当前的时间戳相差超过10秒则判定为过期，则不投递给具体的消费者
        long now = System.currentTimeMillis();
        Map<TopicPartition, List<ConsumerRecord<String, String>>> newRecords = new HashMap<>();
        for (TopicPartition tp : records.partitions()) {
            List<ConsumerRecord<String, String>> tpRecords = records.records(tp);
            List<ConsumerRecord<String, String>> newTpRecords = new ArrayList<>();
            for (ConsumerRecord<String, String> record : tpRecords) {
                if (now - record.timestamp() < EXPIRE_INTERVAL) {
                    newTpRecords.add(record);
                }
            }
            if (!newTpRecords.isEmpty()) {
                newRecords.put(tp, newTpRecords);
            }
        }
        return new ConsumerRecords<>(newRecords);
    }

    @Override
    public void close() {

    }

    /**
     * KafkaConsumer会在提交完消费位移之后调用
     *
     * @param offsets
     */
    @Override
    public void onCommit(Map offsets) {
        offsets.forEach((tp, offset) -> {
            System.out.println(tp + ":" + offset);
        });
    }

    @Override
    public void configure(Map<String, ?> configs) {

    }
}
